#include <cstring>
#include <cstdio>
#include <iostream>
#include <cmath>


#include "Node.h"
#include "TetrahedronState.h"
#include "Tetrahedron.h"
#include "BigVector.h"
#include "FEMModel.h"
#include "Vector.h"
#include "TopologicalOperation.h"
#include "LinearCGState.h"
#include "DeformationState.h"



LinearCGState::LinearCGState ()
{
	residual_len_ = prev_residual_len_ =0;
	iter_count_ = 0;
	residual_tolerance_ = GetNumberSetting ("outer-loop-tolerance");
}

void PStateGradientFunction (void *ps, Real * dest,  Real const*src)
{
	DeformationState *pstate = (DeformationState*) ps;
	pstate->ApplyForceFunction (dest, src);

	BigVectorNegate (dest, dest, pstate->Dimension());
}

void LinearCGState::IterationStep (DeformationState *pstate)
{
	CGIterationStep (pstate->displacement_.UnsafeAccessArray (),
		pstate->fixed_arr_.AccessArray (),
		(void*)pstate,
		&PStateGradientFunction);
}


/*
This code solves

POSITIVE_DEFINITE_OPERATOR x = RHS

where A is a POSITIVE_DEFINITE_OPERATOR.

*/
void LinearCGState::CGIterationStep (Real *displacement,
									bool const *fixed,
									void * optor_info,
									void (*positive_definite_operator)(void*, Real* ,Real const*))
{
	global_iteration_count ++;

	Real beta =  (iter_count_ == 0)
		?  0.0 :  residual_len_ / prev_residual_len_ ;

	iter_count_ ++;

	int n = search_dir_.Size ();
	Real *dir = (Real*)search_dir_.AccessArray ();
	Real *matdir =(Real*) mat_times_search_dir_.AccessArray ();
	Real *residu =(Real*)residual_.AccessArray ();

	BigVectorSaxpy (dir, beta, dir, residu, n);
	BigVectorPartialNullify (dir, dir , fixed, n);
	(*positive_definite_operator) (optor_info, matdir, dir);
	BigVectorPartialNullify (matdir, matdir , fixed, n);   

	Real alpha = residual_len_ / BigVectorIP (matdir, dir, n);

	Real * rp = (Real*)residual_.AccessArray ();

	BigVectorSaxpy (displacement, alpha, dir, displacement, n);
	BigVectorSaxpy (rp, - alpha, matdir, rp, n);  

	/*
	updating residulen is done from a different routine; it also has
	to be done for "correct" solutions */
}



void LinearCGState::UpdateResidual (DeformationState const*pstate)
{
	int n = pstate->Dimension ();

	Real const * delta_f = pstate->delta_external_force_.AccessArray ();
	bool const * fixp =  pstate->fixed_arr_.AccessArray ();

	Real * residu = (Real*)residual_.AccessArray ();
	prev_residual_len_ =residual_len_;
	residual_len_ =0.0;

	bool change = false;
	for (int i = n; i--;)
	{
		if (*delta_f && !*fixp)
		{
			change = true;
			*residu += *delta_f;
		}
		residual_len_ += Sqr (*residu);

		residu ++;
		delta_f ++;
		fixp ++;
	}

	if (isnan (residual_len_) || isinf (residual_len_))
	{
		fprintf (stderr, "Overflow. Exit");
		abort ();
	}

	if (change)
		iter_count_ =0;
}


bool LinearCGState::CGGoodSolution (Real len)const
{
	if (isnan (residual_len_) || isinf (residual_len_))
	{
		fprintf (stderr, "CG Overflow");
		return true;
	}


	bool b = residual_len_ <= len * Sqr (residual_tolerance_);
	return b;
}

/*
It seems nonsensical to keep looking if RESIDUAL_LEN_ is too small,
unfortunately, it makes measuring convergence impossible, since it
introduces a scaling parameter.  */
bool LinearCGState::GoodSolution (DeformationState const * pstate)const
{
	return CGGoodSolution (pstate->external_force_len_sq_);
}


void LinearCGState::PrintNodes () const
{
}


void LinearCGState::Print () const
{
	cout << "res "<<residual_len_ <<endl;
}

void LinearCGState::CGReinit ()
{
	int n = search_dir_.Size ();

	BigVectorNullify (search_dir_.UnsafeAccessArray(),n);
	BigVectorNullify (mat_times_search_dir_.UnsafeAccessArray(),n);
	BigVectorNullify (residual_.UnsafeAccessArray(),n);

	prev_residual_len_ = 0.0;
	iter_count_ =0;
}

/*
Synchronize all state and precomputed information, with the current static
mesh. 
*/
void LinearCGState::Reinit (DeformationState const * pstate) 
{
	CGReinit ();
	residual_len_ = pstate->ComputeResidualForce (residual_.UnsafeAccessArray (),
		pstate->displacement_.AccessArray ());
}


LinkArray< Array<Real> > LinearCGState::VectorVariables ()
{
	Array<Real> *vecs[]
	= {
		&mat_times_search_dir_, &search_dir_, &residual_, 0};

		LinkArray< Array<Real> > vvs;
		for (int i= 0; vecs[i]; i++)
			vvs.Push (vecs[i]);
		return vvs;
}


int LinearCGState::RelaxationCount (FiniteElementModel*fem)const
{
	int count = int (GetUpdateFactor() * 3000.0 / fem->NodeCount ());
	if (count < 1) 
		count = 1;
	return count;
}
extern "C" {
	void
		PrintCGState (LinearCGState*cgs)
	{
		cout << "RL: " << cgs->residual_len_ << " PRL: " << cgs->prev_residual_len_<<endl; 
	}
}
